package com.ruoyi.admin.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.ruoyi.admin.domain.LlmDocumentSegments;
import com.ruoyi.admin.mapper.LlmDocumentsSegmentsMapper;
import com.ruoyi.admin.service.ISearchService;
import com.ruoyi.admin.service.document.ILlmDocumentSegmentsService;
import com.ruoyi.common.constant.MilvusConstants;
import com.ruoyi.common.core.domain.entity.UserInfo;
import com.ruoyi.common.utils.MilvusUtils;
import com.ruoyi.common.utils.StpSystemUtil;
import com.ruoyi.common.utils.StringUtils;
import com.ruoyi.common.utils.embedding.EmbeddingUtils;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
public class SearchServiceImpl implements ISearchService {


    private final ILlmDocumentSegmentsService llmDocumentSegmentsService;

    private final LlmDocumentsSegmentsMapper llmDocumentsSegmentsMapper;

    @Override
    public List<Map<String, Object>> search(String text,
                                            String labelData,
                                            List<String> partitionName,
                                            int topK) {
        UserInfo userInfo = StpSystemUtil.getUserInfo();
        String dataScopeData = userInfo.getDataScopeData();
        if (StringUtils.isBlank(dataScopeData)) {
            return new ArrayList<>();
        }
        List<String> strings = JSONArray.parseArray(dataScopeData, String.class);

        //是否是管理员
        if (StpSystemUtil.isAdmin()) {
            //给空数组就代表是所有权限
            partitionName = new ArrayList<>();
        } else {
            //融合后的权限
            if (!partitionName.isEmpty()) {
                partitionName = strings.stream()
                        .filter(partitionName::contains)
                        .collect(Collectors.toList());
            } else {
                partitionName = strings;
            }
        }


        //在pg库中做标量过滤
        JSONObject jsonObject = JSONObject.parseObject(labelData);
        List<Long> idsLabelData = llmDocumentSegmentsService.getIdsLabelData(jsonObject);
        //创建文本向量
        List<Float> embeddings = EmbeddingUtils.createEmbeddings(text);

        //向量检索
        SearchResp searchResp = MilvusUtils.searchResp(
                MilvusConstants.VECTOR_DOCUMENT_SEGMENT,
                idsLabelData,
                partitionName,
                embeddings,
                topK,
                List.of("id")
        );
        Map<Long, Float> ids = new HashMap<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        searchResults.forEach(searchResult -> {
            searchResult.forEach(searchResult2 -> {
                Map<String, Object> entity = searchResult2.getEntity();
                Long o = (Long) entity.get("id");
                Float score = searchResult2.getScore();
                ids.put(o, score);
            });
        });
        QueryWrapper<LlmDocumentSegments> llmDocumentSegmentsQueryWrapper = new QueryWrapper<>();
        llmDocumentSegmentsQueryWrapper.in(!ids.isEmpty(), "element_id", ids.keySet());

        List<Map<String, Object>> maps = new ArrayList<>();
        List<LlmDocumentSegments> llmDocumentSegments = llmDocumentsSegmentsMapper.selectList(llmDocumentSegmentsQueryWrapper);
        //将查询出来的数据进行封装
        for (LlmDocumentSegments llmDocumentSegment : llmDocumentSegments) {
            Long elementId = llmDocumentSegment.getElementId();
            Map<String, Object> keyObjectHashMap = new HashMap<>();
            keyObjectHashMap.put("elementId", elementId);
            keyObjectHashMap.put("content", llmDocumentSegment.getContent());
            keyObjectHashMap.put("score", ids.get(elementId));
            keyObjectHashMap.put("knowledgeName", llmDocumentSegment.getName());
            maps.add(keyObjectHashMap);
        }

        return maps;
    }
}
